#include <iostream>

#include "visitor.h"

#define TEST

void Visitor::emit() {
    this->backend.emit();
}

void Visitor::setValue(antlr4::tree::ParseTree *node, std::shared_ptr<BaseSymbol> value)
{
    valueTree.put(node, value);
}

std::shared_ptr<BaseSymbol> Visitor::getValue(antlr4::tree::ParseTree *node)
{
    return valueTree.get(node);
}

antlrcpp::Any Visitor::visitUnit(complexCalculatorParser::UnitContext *ctx)
{
#ifdef TEST
    std::cout << "start visit unit" << std::endl;
#endif
    for (auto du : ctx->defUnit())
    {
        visit(du);
    }
    for (auto st : ctx->stat())
    {
        visit(st);
    }
#ifdef TEST
    std::cout << "end visit unit" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitDefUnit(complexCalculatorParser::DefUnitContext *ctx)
{
#ifdef TEST
    std::cout << "start visit def unit" << std::endl;
#endif
    visit(ctx->varUnit());
#ifdef TEST
    std::cout << "end visit def unit" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitVarUnit(complexCalculatorParser::VarUnitContext *ctx)
{
#ifdef TEST
    std::cout << "start visit var def unit" << std::endl;
#endif
    visit(ctx->type());
    for (auto vd : ctx->varDef())
    {
        std::shared_ptr<BaseSymbol> bs(new BaseSymbol());
        bs->setType(getValue(ctx->type())->getType());
        visit(vd);
        bs->setValue(getValue(vd)->getValue());
        bs->setName(getValue(vd)->getName());
        auto addr = this->backend.builder->CreateAlloca(llvm::Type::getInt32Ty(*(this->backend.context)), nullptr);
        bs->setAddr(addr);
        this->backend.builder->CreateStore(bs->getValue(), bs->getAddr());
        this->table.insert(bs);
    }
#ifdef TEST
    std::cout << "end visit var def unit" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitType(complexCalculatorParser::TypeContext *ctx)
{
#ifdef TEST
    std::cout << "start visit type" << std::endl;
#endif
    std::string type_name;
    if (ctx->Int() != nullptr)
    {
        type_name = ctx->Int()->getText();
    }
    if (ctx->Float() != nullptr)
    {
        type_name = ctx->Float()->getText();
    }
    std::shared_ptr<BaseSymbol> bs(new BaseSymbol());
    bs->setType(type_name);
    setValue(ctx, bs);
#ifdef TEST
    std::cout << "end visit type" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitVarDef(complexCalculatorParser::VarDefContext *ctx)
{
#ifdef TEST
    std::cout << "start visit var def" << std::endl;
#endif
    std::shared_ptr<BaseSymbol> bs(new BaseSymbol());
    if (ctx->Identifier() != nullptr)
    {
        // std::cout << ctx->Identifier()->getText() << std::endl;
        bs->setName(ctx->Identifier()->getText());
    }
    if (ctx->initVal() != nullptr)
    {
        visit(ctx->initVal());
        // std::cout << ctx->initVal()->getText() << std::endl;
        bs->setValue(getValue(ctx->initVal())->getValue());
    };
    setValue(ctx, bs);
#ifdef TEST
    std::cout << "end visit var def" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}
antlrcpp::Any Visitor::visitInitVal(complexCalculatorParser::InitValContext *ctx)
{
#ifdef TEST
    std::cout << "start visit init val" << std::endl;
#endif
    visit(ctx->exp());
    setValue(ctx, getValue(ctx->exp()));
#ifdef TEST
    std::cout << "end visit init val" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitExpNumber(complexCalculatorParser::ExpNumberContext *ctx)
{
#ifdef TEST
    std::cout << "start visit exp number" << std::endl;
#endif
    visit(ctx->number());
    setValue(ctx, getValue(ctx->number()));
#ifdef TEST
    std::cout << "end visit exp number" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitNumber(complexCalculatorParser::NumberContext *ctx)
{
#ifdef TEST
    std::cout << "start visit exp number" << std::endl;
#endif
    std::shared_ptr<BaseSymbol> bs(new BaseSymbol());
    if (ctx->INT() != nullptr)
    {
        int x = std::stoi(ctx->INT()->getText());
        auto value = llvm::ConstantInt::get(llvm::Type::getInt32Ty(*(this->backend.context)), x);
        bs->setType("int");
        bs->setValue(value);
    }
    if (ctx->FLOAT() != nullptr)
    {
        double x = std::stod(ctx->FLOAT()->getText());
        auto value = llvm::ConstantFP::get(*(this->backend.context),llvm::APFloat(x));
        bs->setType("float");
        bs->setValue(value);
    }
    setValue(ctx, bs);
#ifdef TEST
    std::cout << "end visit exp number" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitExpAddSub(complexCalculatorParser::ExpAddSubContext *ctx)
{
#ifdef TEST
    std::cout << " start visit exp addsub" << std::endl;
#endif
    std::shared_ptr<BaseSymbol> bs(new BaseSymbol());
    llvm::Value* value;
    visit(ctx->exp(0));
    visit(ctx->exp(1));
    if (getValue(ctx->exp(0))->getType() != getValue(ctx->exp(0))->getType()) {
        std::cout << "wrong type" << std::endl;
        return antlrcpp::Any(nullptr);
    }
    if (ctx->opAddSub()->getText() == "+") {
        if (getValue(ctx->exp(0))->getType() == "int") {
            value = this->backend.builder->CreateAdd(getValue(ctx->exp(0))->getValue(), getValue(ctx->exp(1))->getValue());
        }
        else if (getValue(ctx->exp(0))->getType() == "float") {
            value = this->backend.builder->CreateFAdd(getValue(ctx->exp(0))->getValue(), getValue(ctx->exp(1))->getValue());
        }
        else {
            std::cout << "wrong" << std::endl;
        }
    }
    else if (ctx->opAddSub()->getText() == "-") {
        if (getValue(ctx->exp(0))->getType() == "int") {
            value = this->backend.builder->CreateSub(getValue(ctx->exp(0))->getValue(), getValue(ctx->exp(1))->getValue());
        }
        else if (getValue(ctx->exp(0))->getType() == "float") {
            value = this->backend.builder->CreateFAdd(getValue(ctx->exp(0))->getValue(), getValue(ctx->exp(1))->getValue());
        }
        else {
            std::cout << "wrong" << std::endl;
        }
    }
    else {
        std::cout << "wrong" << std::endl;
    }

    bs->setType(getValue(ctx->exp(0))->getType());
    bs->setValue(value);
    setValue(ctx, bs);
#ifdef TEST
    std::cout << "end visit exp add sub" << std::endl;
#endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitLVal(complexCalculatorParser::LValContext *ctx) {
    #ifdef TEST
    std::cout << " start visit lVal" << std::endl;
    #endif
    std::string name = ctx->Identifier()->getText();
    if (this->table.contains(name)) {
        setValue(ctx, this->table.find(name));
    }
    else {
        std::cout << "do not have symbol" << std::endl;
    }
    // search symbol table 
    #ifdef TEST
    std::cout << " end visit lVal" << std::endl;
    #endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitStatAssign(complexCalculatorParser::StatAssignContext *ctx) {
    #ifdef TEST
    std::cout << " start visit stat assign" << std::endl;
    #endif
    visit(ctx->lVal());
    visit(ctx->exp());
    std::cout << " visit assign value" << std::endl;
    this->backend.builder->CreateStore(getValue(ctx->exp())->getValue(), getValue(ctx->lVal())->getAddr());
    
    getValue(ctx->lVal())->setValue(getValue(ctx->exp())->getValue());
    #ifdef TEST
    std::cout << " end visit stat assign" << std::endl;
    #endif
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitStatExp(complexCalculatorParser::StatExpContext *ctx) {
    std::cout << " visit assign value" << std::endl;
    visit(ctx->exp());
    return antlrcpp::Any(nullptr);
}

antlrcpp::Any Visitor::visitIdentifier(complexCalculatorParser::IdentifierContext *ctx) {
    #ifdef TEST
    std::cout << " start visit identifier" << std::endl;
    #endif
    std::string name = ctx->Identifier()->getText();
    if (this->table.contains(name)) {
        setValue(ctx, this->table.find(name));
    }
    else {
        std::cout << "do not have symbol" << std::endl;
    }
    #ifdef TEST
    std::cout << " end visit identifier" << std::endl;
    #endif
    return antlrcpp::Any(nullptr);
}